package com.tianji.aigc.advisor;

import cn.hutool.core.convert.Convert;
import com.tianji.aigc.enums.AgentTypeEnum;
import com.tianji.aigc.memory.RedisChatMemory;
import org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.api.*;
import org.springframework.ai.chat.model.ChatResponse;
import reactor.core.publisher.Flux;

/**
 * @author wys
 * @version V1.0
 * @date 2025-04-24 18:13
 */
public class RecordOptimizationAdvisor implements CallAroundAdvisor, StreamAroundAdvisor {
    private final RedisChatMemory redisChatMemory;

    public RecordOptimizationAdvisor(RedisChatMemory redisChatMemory) {
        this.redisChatMemory = redisChatMemory;
    }

    /**
     * Around advice that wraps the ChatModel#call(Prompt) method.
     *
     * @param advisedRequest the advised request
     * @param chain          the advisor chain
     * @return the response
     */
    @Override
    public AdvisedResponse aroundCall(AdvisedRequest advisedRequest, CallAroundAdvisorChain chain) {
        AdvisedResponse advisedResponse = chain.nextAroundCall(advisedRequest);
        ChatResponse response = advisedResponse.response();
        if (response == null) {
            return advisedResponse;
        }
        String text = response.getResult().getOutput().getText();
        AgentTypeEnum agentTypeEnum = AgentTypeEnum.agentNameOf(text);
        if (agentTypeEnum != null) {
            String conversationId  = Convert.toStr(advisedResponse.adviseContext()
                    .get(AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY));
            this.redisChatMemory.optimization(conversationId);
        }
        return advisedResponse;
    }

    /**
     * Around advice that wraps the invocation of the advised request.
     *
     * @param advisedRequest the advised request
     * @param chain          the chain of advisors to execute
     * @return the result of the advised request
     */
    @Override
    public Flux<AdvisedResponse> aroundStream(AdvisedRequest advisedRequest, StreamAroundAdvisorChain chain) {
        return chain.nextAroundStream(advisedRequest);
    }

    /**
     * Return the name of the advisor.
     *
     * @return the advisor name.
     */
    @Override
    public String getName() {
        return this.getClass().getSimpleName();
    }

    @Override
    public int getOrder() {
        return Advisor.DEFAULT_CHAT_MEMORY_PRECEDENCE_ORDER - 100;
    }
}
